import requests
import time
from loguru import logger

class ModelWarmer:
    def __init__(self):
        self.api_url = "http://127.0.0.1:5001/chat/stream?q=你好"
        self.timeout = 60  # 请求超时时间（秒）
        
        # 配置日志
        logger.add("warmup.log", rotation="10 MB")
    
    def warmup(self):
        """执行预热请求"""
        try:
            logger.info(f"开始预热模型，请求URL: {self.api_url}")
            start_time = time.time()
            
            # 发送预热请求
            response = requests.get(
                self.api_url,
                stream=True,
                timeout=self.timeout,
                headers={'Accept': 'text/event-stream'}
            )
            
            # 检查响应状态
            if response.status_code == 200:
                logger.info(f"预热请求成功，状态码: {response.status_code}")
                
                # 读取并丢弃响应内容
                for chunk in response.iter_lines():
                    if chunk:
                        # 简单处理响应块，不做实际解析
                        pass
                
                processing_time = time.time() - start_time
                logger.info(f"预热完成，耗时: {processing_time:.2f}秒")
                return True
            else:
                logger.error(f"预热请求失败，状态码: {response.status_code}")
                return False
        except Exception as e:
            logger.error(f"预热请求异常: {str(e)}")
            return False
    
    def wait_for_service(self, max_wait_time=30):
        """等待服务启动"""
        start_time = time.time()
        
        while time.time() - start_time < max_wait_time:
            try:
                # 尝试发送简单请求检查服务是否可用
                response = requests.get("http://127.0.0.1:5001/", timeout=2)
                if response.status_code == 200:
                    logger.info("后端服务已启动，准备进行预热")
                    return True
            except Exception:
                # 服务尚未启动，继续等待
                pass
            
            # 等待1秒后重试
            time.sleep(1)
        
        logger.error(f"在 {max_wait_time} 秒内未检测到服务启动")
        return False

if __name__ == "__main__":
    try:
        warmer = ModelWarmer()
        
        # 等待服务启动
        if warmer.wait_for_service():
            # 执行预热
            success = warmer.warmup()
            if success:
                logger.info("模型预热成功")
            else:
                logger.error("模型预热失败")
        else:
            logger.error("服务未启动，预热中止")
    except Exception as e:
        logger.error(f"预热脚本执行失败: {str(e)}")